import networkx as nx
from typing import Dict, Tuple, List
import heapq

from python.lib import json

from topo1 import build_hybrid_topology_shared_tor
from flow1 import NumpyEncoder
from flow1 import TrafficGenerator #导入类
import numpy as np #将numpy转换成数组形式

def max_flow_first_allocation(
        G: nx.Graph,
        D: Dict[Tuple[str, str], float],
        R: Dict[Tuple[str, str], int],
        optical_capacity: float = 100.0,
        electric_capacity: float = 40.0,
        k_paths: int = 16 #对于流量需求，计算最多k_paths条最短路径
) -> Tuple[Dict, Dict]:
    """
    返回:
        X: 流量分配方案 {(src, dst): allocated}
        remaining_capacity: 链路剩余容量字典
    """
    X = {}
    remaining_capacity = {}

    # # 初始化链路剩余容量
    # for u, v in G.edges():
    #     if R.get((u, v), 0) == 1:
    #         remaining_capacity[(u, v)] = optical_capacity
    #         print(f"{(u,v)}没有重构，所以剩余容量为光链路的剩余容量：{optical_capacity}")
    #     else:
    #         remaining_capacity[(u, v)] = electric_capacity
    #         print(f"{(u, v)}发生了重构，所以剩余容量为电链路的剩余容量：{G[u][v]['capacity']}")

    # 初始化链路剩余容量
    # remaining_capacity = {}
    # for u, v in G.edges():
    #     # if R.get((u, v), 0) == 1:
    #     #     remaining_capacity[(u, v)] = optical_capacity
    #     #     print(f"{(u, v)}没有重构，所以剩余容量为光链路的剩余容量：{optical_capacity}")
    #     # else:
    #     #     remaining_capacity[(u, v)] = electric_capacity
    #     #     print(f"{(u, v)}发生了重构，所以剩余容量为电链路的剩余容量：{G[u][v]['capacity']}")
    #     remaining_capacity [(u, v)] = G[u][v]['capacity']
    #     print(f"{(u, v)}初始化的剩余容量为：{G[u][v]['capacity']}")

    remaining_capacity = {}
    for u, v, data in G.edges(data=True):
        remaining_capacity[(u, v)] = data.get('capacity', 0)
        if not G.is_directed():
            remaining_capacity[(v, u)] = data.get('capacity', 0)

    # 调试信息：打印 remaining_capacity 内容
    print("remaining_capacity 内容：", remaining_capacity)

    sorted_demand = sorted(D.items(), key=lambda x: -x[1])
    print("排序前的通信对：",D)
    print("排序后的通信对：",sorted_demand)

    # 从高需求到低需求依次处理通信对
    for (src, dst), demand in sorted_demand:
        allocated = 0.0 # 当前已经流量分配
        X[(src, dst)] = 0.0  # 是一个字典，例如X[(ToR1, ToR2)] = 56表示当前通信对分配的流量为56

        try:
            # 初始化当前通信对的可行路径为空集
            paths = []
            # 通过最短路径算法找到该通信对的最短路径，找到路径的条数为k_paths
            for path in nx.shortest_simple_paths(G, src, dst, weight='hops'):
                paths.append(path)
                if len(paths) >= k_paths:
                    break
            print(f"从 {src} 到 {dst} 的路径集合：{paths}")
        # 如果找不到最短路径，则跳过当前通信对
        except nx.NetworkXNoPath:
            print(f"找不到从 {src} 到 {dst} 的路径，跳过当前通信对")
            continue

        path_delays = []  # 初始化路径的延迟,是一个集合，包含每条路径的延迟
        # 遍历当前通信对找到的每条路径
        for path in paths:
            print(f"通过最短路径算法找到的路径：{path}")
            # 打印路径的边
            for i in range(len(path) - 1):
                edge = (path[i], path[i + 1])
                print(f"{path}的边：{edge}")

            # 计算路径上所有边的最小剩余容量
            # print("--------------------------------------------------")
            # print(f"路径{path}容量字典中是否包含链路的边")
            # for i in range(len(path) - 1):
            #     if remaining_capacity.get((path[i], path[i + 1])) or remaining_capacity.get((path[i + 1], path[i])):
            #         print(f"路径容量字典中包含{(path[i], path[i + 1])}链路的边")
            #     else :
            #         print(f"路径容量字典中不包含{(path[i], path[i + 1])}链路的边")
            #         break
            # print("路径中所有的边都包含了，可以进行下一步更新最小链路")

            # print("--------------------------------------------------")
            # 设置每条路径的剩余容量为这条路径的每条边剩余容量的最小值
            min_cap = min(max(remaining_capacity.get((path[i], path[i + 1]), remaining_capacity.get((path[i + 1], path[i]))), 0) for i in range(len(path) - 1))
            print(f"{path}的min_cap: {min_cap}")
            # 如果最小剩余容量为0，则跳过该路径
            if min_cap <= 0:
                print(f"{path}的min_cap为0，跳过")
                continue
            # 计算路径的延迟，按照延迟进行排序
            delay = demand / min_cap
            # path_delays.append(delay) # 将延迟添加到路径的延迟列表中
            heapq.heappush(path_delays, (delay, path)) # 将延迟和路径添加到堆中
            # TODO: 这里需要修改，需要计算路径的延迟，然后按照延迟进行排序
            # 输出每条路径的延迟
            print(f"输出路径{path}延迟",delay)  # 为什么所有路径的延时都是一样的？
        # 输出该节点的所有路径延迟
        print("输出路径延迟", path_delays)
        # for j in range(len(path_delays)):
        #     print(path_delays[j][0])
        # 如果堆不为空，则从堆中弹出最小延迟的路径
        while path_delays:
            delay, path = heapq.heappop(path_delays)
            print(f"通信对{(src, dst)}的最小延迟为：{delay},路径为{path}")
            # 计算当前节点选择的最优路径的路径的边的容量（为一个集合）
            # TODO: 这里需要修改，通信节点的剩余链路
            path_caps = [remaining_capacity[(path[i], path[i + 1])] for i in range(len(path) - 1)]
            # 将路径容量的最小值作为路径的容量
            path_cap = min(path_caps)
            # 如果路径容量为0，则跳过该路径
            if path_cap <= 0:
                print(f"{path}的路径容量为0，跳过")
                continue
            # 这里可以添加流量分配等后续操作的代码
            print(f"处理路径 {path}，延迟为 {delay}，路径容量为 {path_cap}")

        # path_delays = [] # 初始化路径的延迟
        # for path in paths:
        #     print(f"通过最短路径算法找到的路径：{path}")
        #     for i in range(len(path) - 1):
        #         edge = (path[i], path[i + 1])
        #         print(f"{path}的边：{edge}")
        #
        #         min_cap = min(remaining_capacity.get(edge, 0)
        #                       for i in range(len(path) - 1))
        #         print(f"{path}的min_cap: {min_cap}")
        #         if min_cap <= 0:
        #             print(f"{path}的min_cap为0，跳过")
        #             continue
        #         delay = demand / min_cap
        #         heapq.heappush(path_delays, (delay, path))
        #
        # print("输出路径延迟",path_delays)
        # while path_delays:
        #     delay, path = heapq.heappop(path_delays)
        #
        #     path_caps = [remaining_capacity[(path[i], path[i + 1])]
        #                  for i in range(len(path) - 1)]
        #     path_cap = min(path_caps)
        #
        #     if path_cap <= 0:
        #         continue
            # 分配的流量大小为路径最小剩余容量和剩余流量的较小值
            allocate = min(demand - allocated, path_cap)
            print(f"通信对{(src, dst)}的最小延迟为：{delay},路径为{path}，分配的流量为{allocate}")

            for i in range(len(path) - 1):
                u, v = path[i], path[i + 1]
                # 更新剩余容量
                # 边的剩余容量为原来边的剩余容量减去分配的流量
                remaining_capacity[(u, v)] -= allocate
            # 更新已经分配的容量
            allocated += allocate
            # 更新通信对分配的流量大小
            X[(src, dst)] = allocated
            # 如果已经分配的容量大于等于需求容量，则跳出循环
            if allocated >= demand - 1e-6:
                break
    # 返回通信对分配的流量大小和剩余容量
    return X, remaining_capacity  # 返回两个值


# 测试用例
if __name__ == "__main__":
    # G = nx.Graph()
    # edges = [
    #     ('A', 'B', {'hops': 1}),
    #     ('B', 'C', {'hops': 1}),
    #     ('A', 'C', {'hops': 2}),
    #     ('C', 'D', {'hops': 1})
    # ]
    # G.add_edges_from(edges)
    #
    # R = {('A', 'B'): 1}
    #
    # D = {
    #     ('A', 'D'): 80.0,
    #     ('B', 'C'): 50.0,
    #     ('A', 'C'): 30.0
    # }

    G = build_hybrid_topology_shared_tor(num_tor=4)

    # 接受返回的流量矩阵为Ring-AllReduce
    generator = TrafficGenerator(num_nodes=4) #创建实例
    np_matrix = generator.generate_all_to_all()  # numpy数组

    # print("=========从json文件中收集返回的all-to-all流量数据集===================")
    # with open('traffic_simulation.json', 'r') as f:
    #     data = json.load(f)
    #     # 输出时间为0.12时候的矩阵
    #     for record in data:
    #         if record['timestamp'] == 0.12:
    #             matrix = record['matrix']
    #             # 转换成nmpy数组
    #             np_matrix = np.array(matrix)
    #             print(np_matrix)
    # 转换为字典格式
    D = {}
    for i in range(np_matrix.shape[0]):
        for j in range(np_matrix.shape[1]):
            if i != j:
                src_tor = f"ToR{i + 1}"
                dst_tor = f"ToR{j + 1}"
                D[(src_tor, dst_tor)] = np_matrix[i, j]
    # 初始化链路类型标记字典R（示例）
    R = {('ToR1', 'Optical_SW1'): 0}  # 1表示光链路
    print("需要重构的节点对为：",R)
    # 接收两个返回值
    allocation, remaining_cap = max_flow_first_allocation(G, D, R)
    # for (src, dst), demand in D.items():
    #     print(f"{src}->{dst}: 申请{demand}G 分配{allocation.get((src, dst), 0):.1f}G")


    # print("流量分配结果:")
    # for (src, dst), demand in D.items():
    #     print(f"{src}->{dst}: 申请{demand}G 分配{allocation.get((src, dst), 0):.1f}G")
    #
    # print("\n最终链路剩余容量:")
    # for u, v in G.edges():
    #     print(f"{u}-{v}: {remaining_cap.get((u, v), 0):.1f}G")
    #
    # print("当前图结构节点列表:", list(G.nodes()))
    # print("图结构边列表:", list(G.edges()))

    # 构建结果字典
    result = {
        "topology": {
            "nodes": list(G.nodes()),
            "edges": [
                {
                    "source": u,
                    "target": v,
                    "type": G[u][v]['type'],
                    "initial_capacity": G[u][v].get('capacity', 0),
                    "remaining_capacity": remaining_cap.get((u, v), 0)
                }
                for u, v in G.edges()
            ]
        },
        "allocation": [
            {
                "source": src,
                "destination": dst,
                "demand": D[(src, dst)],
                "allocated": allocation.get((src, dst), 0)
            }
            for (src, dst) in D
        ]
    }

    # 保存为JSON文件
    with open("allocation_results.json", "w") as f:
        json.dump(result, f, indent=2, cls=NumpyEncoder)

    print("结果已保存至 allocation_results.json")

    # 绘制流量分配之后的新流量矩阵图
    # TODO 流量分配之后的流量矩阵图与之前的相同
    import matplotlib.pyplot as plt

    hotspot_nodes = [0, 1]
    hotspot_traffic_over_time = {node: [] for node in hotspot_nodes}
    timestamps = []

    for record in allocation:
        print("record 的内容：", record)

        try:
            record_dict = dict(record)
        except ValueError as e:
            print(f"转换为字典时出错：{e}")
        # 将元组转化为字典
        record_dict = dict(record)
        matrix = np.array(record_dict["matrix"])
        timestamp = record["timestamp"]
        timestamps.append(timestamp)
        for node in hotspot_nodes:
            hotspot_traffic_over_time[node].append(matrix[node, :].sum())

    plt.figure(figsize=(10, 6))
    for node in hotspot_nodes:
        plt.plot(timestamps, hotspot_traffic_over_time[node], label=f'Node {node}')

    plt.title('Hotspot Node Traffic over Time')
    plt.xlabel('Time (s)')
    plt.ylabel('Traffic Volume')
    plt.legend()
    plt.show()



